import os, sqlite3, re
from src.core import log
from threading import Lock
from functools import wraps
from sqlalchemy import and_, or_
from sqlalchemy.pool import QueuePool
from sqlalchemy.exc import SQLAlchemyError
from src.core.variables.exception import *
from src.core.variables.function import SQL
from sqlalchemy import create_engine, MetaData
from src.core.variables.constant import DATABASE
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, scoped_session
from sqlalchemy.orm.attributes import InstrumentedAttribute


def auto_change_table_structure(func):
    """
    定义的表结构与数据库表结构发生变化：
        新增字段时：自动创建
        缺少字段时 和 字段更名时：直接退出并提示手动修改数据库
    """

    @wraps(func)
    def wapper(self, *args, **kwargs):
        if DATABASE.DATABASE_TYPE == "sqlite":
            db_obj = kwargs.get("db_obj")
            db_path = kwargs.get("db_path")
            if os.path.exists(db_path):
                with sqlite3.connect(db_path, timeout=10) as conn:
                    cursor = conn.cursor()
                    for table in cursor.execute(SQL.ALL_TABLES()).fetchall():
                        field_list = []
                        new_field_list = []
                        for field in cursor.execute(
                            SQL.ALL_FILEDS(table[0])
                        ).fetchall():
                            field_list.append(field[1])
                        for index, value in getattr(
                            db_obj,
                            re.sub(r"^\s*\w", lambda m: m.group().upper(), table[0]),
                        ).__dict__.items():
                            if (
                                type(value) is InstrumentedAttribute
                                and index not in field_list
                            ):
                                new_field_list.append(index)
                        for i in new_field_list:
                            field_obj = getattr(getattr(db_obj, table[0].capitalize()), i)
                            conn.execute(
                                SQL.ALTER_FIELD(field_obj.table.name, field_obj.key)
                            )
                            conn.commit()
        elif DATABASE.DATABASE_TYPE == "mysql":
            log.CRITICAL(
                "当前版本暂不支持mysql数据库，请使用sqlite数据库！请修改配置文件中的db.type为sqlite"
            )
        result = func(self, *args, **kwargs)
        return result

    return wapper


class SQLiteMultiDB:
    def __init__(self):
        self.engines = {}
        self._session_factories = {}
        self._lock = Lock()
        self._metadata = MetaData()  # 共享元数据
        self._model_bindings = {}  # 存储模型与数据库的绑定关系
        self._db_metadata_map = {}  # 数据库到其实际使用元数据的映射
        self._registered_tables = set()  # 记录已注册的表名
        self.and_ = and_
        self.or_ = or_

    @auto_change_table_structure
    def add_database(self, db_name, db_path, db_obj):
        """添加新数据库配置"""
        with self._lock:
            if db_name in self.engines:
                raise ValueError(f"数据库 '{db_name}' 已存在")
            os.makedirs(os.path.dirname(db_path), exist_ok=True)
            # 创建数据库引擎
            engine = create_engine(f"sqlite:///{db_path}", poolclass=QueuePool, max_overflow=50)
            # 初始化WAL模式
            # with engine.connect() as conn:
            #     conn.execute(text("PRAGMA journal_mode=WAL"))
            self.engines[db_name] = engine
            self._session_factories[db_name] = scoped_session(
                sessionmaker(bind=engine, autocommit=False, autoflush=False)
            )
            # 为该数据库创建独立的元数据副本
            self._db_metadata_map[db_name] = MetaData()
            log.DEBUG(f"数据库 '{db_name}' 创建成功，路径: {db_path}")
        # 注册模型
        model_count = 0
        if hasattr(db_obj, "__dict__"):  # 确保db_obj是可迭代的
            for name, obj in db_obj.__dict__.items():
                if isinstance(obj, type) and issubclass(obj, Base):
                    with self._lock:
                        if db_name not in self.engines:
                            raise KeyError(f"数据库 '{db_name}' 未找到")
                        # 检查表是否已注册
                        table_name = obj.__tablename__
                        if (db_name, table_name) in self._registered_tables:
                            continue  # 跳过已注册的表
                        # 将模型绑定到共享元数据
                        if not hasattr(obj, "metadata"):
                            obj.metadata = self._metadata
                        # 将表定义复制到该数据库的独立元数据中
                        table = obj.__table__
                        table_copy = table.tometadata(self._db_metadata_map[db_name])
                        # 记录模型与数据库的绑定关系
                        self._model_bindings[obj.__name__] = db_name
                        self._registered_tables.add((db_name, table_name))
                        model_count += 1
                        log.DEBUG(f"数据库 '{db_name}' 注册表 '{table_name}' 成功")
        if model_count == 0:
            # 即使没有表也确保元数据存在
            self._db_metadata_map.setdefault(db_name, MetaData())

    def create_tables(self, db_name=None):
        """为指定数据库创建所有注册的表"""
        with self._lock:
            if db_name:
                # 创建指定数据库的所有表
                if db_name not in self.engines:
                    raise KeyError(f"数据库 '{db_name}' 未找到")
                # 确保元数据存在
                metadata = self._db_metadata_map.get(db_name, MetaData())
                if not metadata.tables:
                    log.CRITICAL(f"数据库 '{db_name}' 没有可创建的表")
                    return
                metadata.create_all(bind=self.engines[db_name])
                # 记录创建的表
                tables = [t.name for t in metadata.tables.values()]
                log.DEBUG(
                    f"数据库 '{db_name}' 表创建完成: {', '.join(tables) if tables else '没有创建任何表'}"
                )
            else:
                # 为所有数据库创建各自的表
                total_tables = 0
                created_dbs = []
                for db_name, engine in self.engines.items():
                    metadata = self._db_metadata_map.get(db_name, MetaData())
                    if metadata.tables:
                        metadata.create_all(bind=engine)
                        tables = [t.name for t in metadata.tables.values()]
                        table_count = len(tables)
                        total_tables += table_count
                        log.DEBUG(f"数据库 '{db_name}' 表创建完成: {', '.join(tables)}")
                        created_dbs.append(db_name)
                    else:
                        log.CRITICAL(f"数据库 '{db_name}' 没有可创建的表")
                if created_dbs:
                    log.DEBUG(
                        f"所有表创建完成，涉及数据库: {', '.join(created_dbs)}，共创建 {total_tables} 张表"
                    )
                else:
                    log.CRITICAL("没有在任何数据库中创建表")

    def get(self, db_name) -> scoped_session:
        """获取指定数据库的会话"""
        with self._lock:
            if db_name not in self._session_factories:
                raise KeyError(f"数据库 '{db_name}' 未找到")
            return self._session_factories[db_name]()

    def close_all(self):
        """关闭所有连接"""
        with self._lock:
            for name, factory in self._session_factories.items():
                factory.remove()
                self.engines[name].dispose()
            log.DEBUG("所有数据库连接已关闭")

    def commit(self, db_session):
        """提交会话"""
        try:
            db_session.commit()
            return True
        except SQLAlchemyError as e:
            log.ERROR(f"Commit 失败: {e}")
        except Exception as e:
            db_session.rollback()
            log.ERROR(f"提交数据库时出现未知错误: {e}")
        return False


db_manager = SQLiteMultiDB()
Base = declarative_base(metadata=db_manager._metadata)
